import cv2
import numpy as np
from numba import cuda
import time
import math


# gpu function
@cuda.jit
def process_gpu(img, channels):
    tx = cuda.blockIdx.x * cuda.blockDim.x + cuda.threadIdx.x
    ty = cuda.blockIdx.y * cuda.blockDim.y + cuda.blockIdx.y
    for c in range(channels):
        color = img[tx, ty, c] * 2.0 + 30
        if color > 255:
            img[tx, ty, c] = 255
        elif color < 0:
            img[tx, ty, c] = 0
        else:
            img[tx, ty, c] = color


# cpu function
def process_cpu(img, dst):
    rows, cols, channels = img.shape
    for i in range(rows):
        for j in range(cols):
            for c in range(channels):
                color = img[i, j][c] * 2.0 + 30
                if color > 255:
                    dst[i, j][c] = 255
                elif color < 0:
                    dst[i, j][c] = 0
                else:
                    dst[i, j][c] = color


if __name__ == "__main__":
    # create image
    img = cv2.imread('image_deal.jpg')
    rows, cols, channels = img.shape
    dst_cpu = img.copy()
    dst_gpu = img.copy()
    start_cpu = time.time()
    process_cpu(img, dst_cpu)
    end_cpu = time.time()
    time_cpu = (end_cpu - start_cpu)
    print("CPU process time: " + str(time_cpu))

    # GPU function
    dImg = cuda.to_device(img)
    threads_per_block1 = (16, 16)
    blocks1_per_grid_x = int(math.ceil(rows / threads_per_block1[0]))
    blocks1_per_grid_y = int(math.ceil(cols / threads_per_block1[1]))
    blocks_per_grid = (blocks1_per_grid_x, blocks1_per_grid_y)
    cuda.synchronize()
    start_gpu1 = time.time()
    process_gpu[blocks_per_grid, threads_per_block1](dImg, channels)
    cuda.synchronize()
    end_gpu1 = time.time()
    dst_gpu1 = dImg.copy_to_host()
    time_gpu1 = (end_gpu1 - start_gpu1)

    # save
    # cv2.imwrite("result_cpu.jpg", dst_cpu)
    cv2.imwrite("result_gpu1.jpg", dst_gpu1)

    print("GPU1 process time: " + str(time_gpu1))

    # GPU function
    dImg = cuda.to_device(img)
    threads_per_block2 = (32, 32)
    blocks2_per_grid_x = int(math.ceil(rows / threads_per_block2[0]))
    blocks2_per_grid_y = int(math.ceil(cols / threads_per_block2[1]))
    blocks2_per_grid = (blocks2_per_grid_x, blocks2_per_grid_y)
    cuda.synchronize()
    start_gpu2 = time.time()
    process_gpu[blocks2_per_grid, threads_per_block2](dImg, channels)
    cuda.synchronize()
    end_gpu2 = time.time()
    dst_gpu2 = dImg.copy_to_host()
    time_gpu2 = (end_gpu2 - start_gpu2)

    # save
    cv2.imwrite("result_gpu2.jpg", dst_gpu2)

    print("GPU2 process time: " + str(time_gpu2))
